import cv2
import gradio as gr
from ultralytics import YOLO


def detect(input):
    print(input)
    model = YOLO(r"C:\Users\angel\Desktop\test\model\best.pt")
    metrics = model.predict([input])
    image = metrics[0].plot()
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image


with gr.Blocks() as demo:
    input = gr.Image(type="filepath", label="输入图片")
    output = gr.Image(type="numpy", label="输出图片")
    detect_btn = gr.Button("开始检测")
    detect_btn.click(fn=detect, inputs=input, outputs=output)

demo.launch()
